from osgeo import gdal,ogr,osr,gdalconst
from itertools import product
from numba import jit
from tqdm import tqdm
import numpy as np
import os


@jit
def my_tqdm(iterator):
    return tqdm(iterator)


@jit(nopython=True)  # 使用Numba的JIT编译器来加速函数
def connected8_2_connected4(in_arr):
    pattern_1 = np.array([[255, 0], [0, 255]], dtype=np.float32)
    pattern_2 = np.array([[0, 255], [255, 0]], dtype=np.float32)
    replacement = np.array([[255, 255], [255, 255]], dtype=np.float32)
    out_arr = in_arr.copy()
    rows, cols = in_arr.shape # 行，列
    # for i,j in product(range(rows-1),range(cols-1)):
    for i in range(rows-1):
        for j in range(cols-1):
            tmp_arr = out_arr[i:i+2,j:j+2].copy() 
            if np.array_equal(tmp_arr, pattern_1):
                out_arr[i:i+2,j:j+2] = replacement
            elif np.array_equal(tmp_arr, pattern_2):
                out_arr[i:i+2,j:j+2] = replacement
    if not np.array_equal(in_arr,out_arr):
        out_arr = connected8_2_connected4(out_arr)
    return out_arr


@jit(nopython=True)  # 使用Numba的JIT编译器来加速函数
def del_single_pixel(in_arr):
    pattern = np.array([[0, 0, 0], [0, 255, 0], [0, 0, 0]], dtype=np.uint8)
    replacement = np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]], dtype=np.uint8)
    out_arr = in_arr.copy()
    rows, cols = in_arr.shape # 行，列
    # for i,j in product(range(rows-2),range(cols-2)):
    for i in range(rows-2):
        for j in range(cols-2):
            tmp_arr = out_arr[i:i+3,j:j+3].copy() 
            if np.array_equal(tmp_arr, pattern):
                out_arr[i:i+3,j:j+3] = replacement
    return out_arr


def raster_2_geometry(tif_path):
    dataset = gdal.Open(tif_path)
    projection = dataset.GetProjection()
    srs = osr.SpatialReference(wkt=projection)
    # if srs.IsGeographic():print("使用的是地理坐标系")
    epsg_code = srs.GetAttrValue('AUTHORITY', 1)
    if epsg_code != '4326':
        mem_path = '/vsimem/output.tif'
        print("转WGS84坐标系")
        gdal.Warp(
            mem_path, 
            dataset, 
            dstSRS='EPSG:4326',
            resampleAlg=gdal.GRA_Bilinear,
            callback=gdal.TermProgress_nocb)
    else:
        print("该栅格使用的是WGS84坐标系。")
        mem_path = tif_path
    del dataset
    
    dataset = gdal.Open(mem_path) 
    xsize = dataset.RasterXSize
    ysize = dataset.RasterYSize
    # band_count = dataset.RasterCount
    geotransform = dataset.GetGeoTransform()
    projection = dataset.GetProjection()
    band = dataset.GetRasterBand(1)
    data_type = band.DataType
    # nodata = band.GetNoDataValue()
    mask_band = band.GetMaskBand()
    array = mask_band.ReadAsArray()
    del dataset 
    
    array = connected8_2_connected4(array)
    array = del_single_pixel(array)
    
    driver = gdal.GetDriverByName("MEM") # 'GTiff'
    mem_ds = driver.Create("", xsize, ysize, 1, data_type)
    mem_ds.SetGeoTransform(geotransform)  # 地理变换信息
    mem_ds.SetProjection(projection)  # 设置投影
    band = mem_ds.GetRasterBand(1)
    band.WriteArray(array)
    band.SetNoDataValue(0)
    mask_band = band.GetMaskBand()
    mem_ds.FlushCache()
 
    srs=osr.SpatialReference()
    srs.ImportFromEPSG(4326)
    drivername = "Memory" # "ESRI Shapefile" # "MEM"
    # shp_path = r"/data/fengyy/dateset/temp/202403/131123_武强县20230726_test4.shp"
    driver = ogr.GetDriverByName(drivername) # 'ESRI Shapefile'
    # if os.path.exists(shp_path):driver.DeleteDataSource(shp_path)
    vector_ds = driver.CreateDataSource('')
    layer = vector_ds.CreateLayer('layer',srs,geom_type=ogr.wkbPolygon)
    # layer.CreateField(ogr.FieldDefn("ID", ogr.OFTInteger))
    print("栅格转矢量")
    gdal.Polygonize(
        band, 
        mask_band, 
        layer, 
        -1, 
        options=[], 
        callback=gdal.TermProgress_nocb)
    del mem_ds
    featureCount = layer.GetFeatureCount()
    if featureCount == 1:
        feature = layer.GetFeature()
        reslut = geometry = feature.geometry()
    else:
        multi_polygon = ogr.Geometry(ogr.wkbMultiPolygon)
        for index in range(featureCount):
            feature = layer.GetFeature(index)
            geometry = feature.geometry()
            multi_polygon.AddGeometry(geometry)
        reslut = multi_polygon.ExportToWkt()
    vector_ds.Destroy()
    print(reslut)
    return reslut


def projection_2_spatil(tif_path):
    # 打开栅格数据集
    dataset = gdal.Open(tif_path)
    # 获取WKT格式的空间参考信息
    wkt = dataset.GetProjection()
    # 创建空间参考对象
    spatial_ref = osr.SpatialReference(wkt=wkt)
    # 检查空间参考是否存在
    if wkt == '':
        raise "无空间参考信息"
    else:
        # 检查是否为地理坐标系
        if spatial_ref.IsGeographic():
            print("使用的是地理坐标系")
            # 获取EPSG代码
            epsg_code = spatial_ref.GetAttrValue('AUTHORITY', 1)
            # 检查是否是WGS84
            if epsg_code == '4326':
                print("该栅格使用的是WGS84坐标系。")
                # vrt_path = re.sub("\..*$",".vrt",tif_path)
                mem_path = '/vsimem/output.tif'
                gdal.Warp(mem_path, dataset, dstSRS='EPSG:4326')
                return mem_path
            else:
                # vrt_path = re.sub("\..*$",".vrt",tif_path)
                mem_path = '/vsimem/output.tif'
                # 使用gdal.Warp()进行重新投影
                gdal.Warp(mem_path, dataset, dstSRS='EPSG:4326')
                return mem_path
        # 检查是否为投影坐标系
        elif spatial_ref.IsProjected():
            print("使用的是投影坐标系")
            # 使用gdal.Warp()进行重新投影
            # vrt_path = re.sub("\..*$",".vrt",tif_path)
            mem_path = '/vsimem/output.tif'
            gdal.Warp(mem_path, dataset, dstSRS='EPSG:4326')
            return mem_path


def main(tif_path,shp_path):
    # 打开栅格数据集
    raster_path = projection_2_spatil(tif_path)
    dataset = gdal.Open(raster_path,gdalconst.GA_Update)
    transform = dataset.GetGeoTransform()
    columns, rows = dataset.RasterXSize, dataset.RasterYSize
    min_area = abs(transform[1]*transform[5])*10
    wkt_projection = dataset.GetProjection()
    band = dataset.GetRasterBand(2)
    nodata = band.GetNoDataValue()
    band_array = band.ReadAsArray()
    # mask_band = band.GetMaskBand()
    # array = mask_band.ReadAsArray()
    mask_array = np.zeros_like(band_array, dtype=np.float32)
    mask_array[band_array != nodata] = 1
   
    print(mask_array[mask_array == 1])
    result_array = connected8_2_connected4(mask_array)
    # array = band.ReadAsArray()
    # mask_band = band.GetMaskBand()
    # mask_array = mask_band.ReadAsArray()
    
    
    driver = gdal.GetDriverByName('MEM')
    mem_ds = driver.Create('', columns, rows, 1, gdal.GDT_Float32)
 
    # 设置栅格的一些元数据
    mem_ds.SetGeoTransform(transform)  # 地理变换信息
    mem_ds.SetProjection(wkt_projection)  # 设置投影
    
    print(result_array[result_array == 1])
    # 创建数据并写入内存栅格数据集
    data = np.random.rand(rows, columns)
    band = mem_ds.GetRasterBand(1)
    band.WriteArray(result_array)
    band.SetNoDataValue(nodata)
    band = mem_ds.GetRasterBand(1)
    mask_band = band.GetMaskBand()
    
    # band.WriteArray(result_array)
    # band = dataset.GetRasterBand(1)
    srs = osr.SpatialReference()
    srs.ImportFromWkt(wkt_projection)
    # 创建输出矢量数据源
    driver = ogr.GetDriverByName('ESRI Shapefile') # Memory ESRI Shapefile
    mem_ds = driver.CreateDataSource(shp_path) # memory
    layer = mem_ds.CreateLayer('polygon_layer', srs, ogr.wkbPolygon)
    # 添加字段
    fieldDef = ogr.FieldDefn('bandvalue', ogr.OFTInteger)
    layer.CreateField(fieldDef)
    options = []
    # options.append('8CONNECTED=8')  # 使用 8 连通区域生长算法
    gdal.Polygonize(
        band, 
        None, #mask_band, 
        layer, 
        -1, 
        options=options, 
        callback=gdal.TermProgress_nocb
    )
    geometrys = []
    for index in range(layer.GetFeatureCount()):
        fearture = layer.GetFeature(index) 
        geometry = fearture.geometry()
        geometrys.append(geometry.ExportToWkt())
        area = geometry.GetArea()
        print(area)
        # print(area)
        # if min_area <= area:
        #     geometrys.append(geometry)
        # else:
        #     layer.DeleteFeature(index)
    # print(geometrys)
    # for (polygonwkt1, wkt2) in combinations(geometrys, 2):
    #     print(polygon1)
        # print(polygon1.Touches(polygon2))
        
    # geometry = geometrys[areas.index(max(areas))]
    # geometry_wkt = geometry.ExportToWkt()
    mem_ds.FlushCache()
    mem_ds.Destroy()
    del dataset
    return


tif_path = r"/data/fengyy/dateset/temp/202403/131123_武强县20230726.tif"
shp_path = r"/data/fengyy/dateset/temp/202403/131123_武强县20230726-01.shp"
# main(tif_path,shp_path)
mem_path = r"/data/fengyy/dateset/temp/202403/131123_武强县20230726_test1.tif"
# tif_2_mask_mem(tif_path,mem_path)
# raster_path = r"/data/fengyy/dateset/temp/202403/131123_武强县20230726_test.tif"
vector_path = r"/data/fengyy/dateset/temp/202403/131123_武强县20230726_test.shp"
# drivername = "ESRI Shapefile"

tif_2_shp(tif_path)



